package com.lagou.homework.sparkstream

import java.util

import org.apache.kafka.common.TopicPartition
import org.apache.spark.streaming.kafka010.OffsetRange
import redis.clients.jedis.{Jedis, JedisPool, JedisPoolConfig}

import scala.collection.mutable

object OffsetRedisUtils {
    private val config = new JedisPoolConfig
    private val redisHost = "linux121"
    private val redisPort = 6379
    //最大连接
    config.setMaxTotal(10)
    //最大空闲
    config.setMaxIdle(5)

    private val pool = new JedisPool(config, redisHost, redisPort, 10000)

    private def getRedisConnection: Jedis = pool.getResource

    private val topicPrefix = "kafka:topic"

    //key的格式 => prefix:topic:groupId
    private def getKey(topic: String, groupId: String): String = s"$topicPrefix:$topic:$groupId"

    //从Jedis中获取offsets
    def getOffsetsFromRedis(topics: Array[String], groupId: String): Map[TopicPartition, Long] = {
        val jedis: Jedis = getRedisConnection
        val offsets: Array[mutable.Map[TopicPartition, Long]] = topics.map { topic =>
            import scala.collection.JavaConverters._

            jedis.hgetAll(getKey(topic, groupId))
                    .asScala
                    .map { case (partition, offset) =>
                        new TopicPartition(topic, partition.toInt) -> offset.toLong
                    }
        }
        //归还资源
        jedis.close()
        offsets.flatten.toMap
    }

    //将offsets保存到redis
    def saveOffsetsToRedis(ranges: Array[OffsetRange], groupId: String): Unit = {
        val jedis: Jedis = getRedisConnection

        ranges.map(range => (range.topic, range.partition.toString -> range.untilOffset.toString))
                .groupBy(_._1)
                .foreach { case (topic, buffer) =>
                    import scala.collection.JavaConverters._
                    val maps: util.Map[String, String] = buffer.map(_._2).toMap.asJava
                    jedis.hmset(getKey(topic, groupId), maps)
                }

//        println("saveOffsets")
        //关闭资源
        jedis.close()
    }
}
